# Ultralytics YOLO 🚀, GPL-3.0 license
import torch
from typing import Any, Dict
from torch.utils.data import DataLoader
from torchvision import transforms
from ultralytics.yolo.engine.validator import BaseValidator
from ultralytics.yolo.v8.detect import DetectionValidator
from ultralytics.yolo.utils import DEFAULT_CFG, ops
from ultralytics.yolo.utils.metrics import ConfusionMatrix, DetMetrics

from .data import SOD, sod_collate
from .process import preprocess_batch, postprocess_batch


class SODValidator(DetectionValidator):
    def __init__(self, dataloader=None, save_dir=None, pbar=None, logger=None, args=None):
        self.ch = 5
        self.args = args
        if dataloader is None:
            self.dataloader = dataloader = self.get_dataloader(self.get_dataset(args.data), batch_size = args.batch)
        BaseValidator.__init__(self, dataloader, save_dir, pbar, logger, args)
        self.args.task = 'detect'
        self.data_dict = self.args.data
        self.is_coco = False
        self.class_map = None
        self.metrics = DetMetrics(save_dir=self.save_dir)
        self.iouv = torch.linspace(0.5, 0.95, 10)  # iou vector for mAP@0.5:0.95
        self.niou = self.iouv.numel()

    def get_dataset(self, data: Dict[str, Any]):
        test_dataset = None
        padding = lambda x, y, d: transforms.Resize((d * (x // d), d * (y // d)), antialias = True)
        seq_len = data["seq_len"] if "seq_len" in data else 10
        if data["name"] == "pku_davis_sod":
            test_dataset = SOD(
                data["path"],
                "test",
                transform = padding(240, 320, 64),
                seq_len = seq_len,
                augment = False
            )
            data["nc"] = len(test_dataset.classes.keys())
            data["names"] = test_dataset.classes
        else:
            return NotImplementedError("Unknown dataset %s" % (data["name"],))
        return test_dataset

    def get_dataloader(self, dataset_path, batch_size, mode="train", rank=0):
        # TODO: manage splits differently
        # calculate stride - check if model is initialized
        return DataLoader(
            dataset_path,
            batch_size = batch_size,
            num_workers = self.args.workers,
            collate_fn = sod_collate,
            drop_last = True
        )

    def preprocess(self, batch):
        batch = preprocess_batch(batch, self.device, torch.half if self.args.half else torch.float)

        for k in ["batch_idx", "cls", "bboxes"]:
            batch[k] = batch[k].to(self.device)

        nb = len(batch["img"])
        self.lb = [torch.cat([batch["cls"], batch["bboxes"]], dim=-1)[batch["batch_idx"] == i]
                   for i in range(nb)] if self.args.save_hybrid else []  # for autolabelling

        return batch

    def init_metrics(self, model):
        head = model.model[-1] if self.training else model.model.model[-1]
        self.class_map = list(range(1000))
        self.args.save_json |= self.is_coco and not self.training  # run on final val if training COCO
        self.nc = head.nc
        self.names = model.names
        self.metrics.names = self.names
        self.metrics.plot = self.args.plots
        self.confusion_matrix = ConfusionMatrix(nc=self.nc)
        self.seen = 0
        self.jdict = []
        self.stats = []

    def postprocess(self, preds, batch = None):
        preds = ops.non_max_suppression(preds,
                                        self.args.conf,
                                        self.args.iou,
                                        labels=self.lb,
                                        multi_label=True,
                                        agnostic=self.args.single_cls,
                                        max_det=self.args.max_det)
        if batch is not None:
            batch = postprocess_batch(batch)
        return preds


def val(cfg=DEFAULT_CFG, use_python=False):
    model = cfg.model or "yolov8n.pt"
    data = cfg.data or "coco128.yaml"

    args = dict(model=model, data=data)
    if use_python:
        from ultralytics import YOLO
        YOLO(model).val(**args)
    else:
        validator = SODValidator(args=args)
        validator(model=args['model'])


if __name__ == "__main__":
    val()
